"""线性回归量化模型"""
import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from typing import Dict, Optional
from .base_model import BaseQuantModel


class LinearRegressionModel(BaseQuantModel):
    """线性回归模型（回归：未来收益）"""
    
    def __init__(self, **kwargs):
        """
        初始化线性回归模型
        
        Args:
            **kwargs: LinearRegression参数
        """
        super().__init__(model_name="LinearRegression", task_type="regression")
        self.model_params = kwargs if kwargs else {}
    
    def _build_model(self, **kwargs):
        """构建线性回归模型"""
        params = {**self.model_params, **kwargs}
        self.model = LinearRegression(**params)
    
    def _train_model(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: Optional[pd.DataFrame] = None,
        y_val: Optional[pd.Series] = None,
        **kwargs
    ) -> Dict:
        """训练线性回归模型"""
        self.model.fit(X_train, y_train)
        
        # 计算训练集R²
        train_score = self.model.score(X_train, y_train)
        
        result = {
            'train_r2': train_score
        }
        
        # 如果有验证集，计算验证集R²
        if X_val is not None and y_val is not None:
            val_score = self.model.score(X_val, y_val)
            result['val_r2'] = val_score
        
        # 计算训练集MSE
        train_pred = self.model.predict(X_train)
        train_mse = np.mean((train_pred - y_train) ** 2)
        result['train_mse'] = train_mse
        
        if X_val is not None and y_val is not None:
            val_pred = self.model.predict(X_val)
            val_mse = np.mean((val_pred - y_val) ** 2)
            result['val_mse'] = val_mse
        
        return result
    
    def _predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        """预测值"""
        return self.model.predict(X)

